#include "samplers.fx"

static const float pi = 3.14159265359;

//_____S T R U C T U R E S_________________________________
struct FireParticle
{
	float4 position		: POSITION;
	float4 direction	: DIRECTION; 
	float4	params		: PARAMETERS;	// emitter : .X: ID, .Y: active: Z: age .W: unique ID (0,1,2,...)
										// particle: .X: age .Y: alpha   Z:     .W: type (=-1)
};

struct FireParticle_ID_VS
{
	float4 position		: POSITION;
	float4 direction	: DIRECTION; 
	float4	params		: PARAMETERS;	// emitter : .X: ID, .Y: active: Z: age .W: type
										// particle: .X: age .Y: alpha   Z:     .W: type
	uint	vertexId	:  SV_VertexID;
};

struct FireParticle_ID_GS
{
	float4 position		: POSITION;
	float4 direction	: DIRECTION; 
	float4	params		: PARAMETERS;	// emitter : .X: ID, .Y: active: Z: age .W: type
										// particle: .X: age .Y: alpha   Z:     .W: type
	uint	vertexId	:  VertexID;
	float3	emitterWS	:	EMITTER_WS;
};

struct GS_OUT
{
	float4	position	:	SV_POSITION;
};

//________________V A R I A B L E S___________________________________

float4x4	matrices[50];
float		dt;
float		fHeightThreshold;
Texture2D	tex_heat;
Texture2D	tex_noise;
int			activate_emitter;
int			heatTex_dimension;
float		heat_threshold;
float		time;
float		fVal;
uint		iMaxEmitters;
bool		bUseDiskDistribution;
float3		vWindDir;
float		fWindPower;
bool		bActive;
float		fParticleAge;

//________________H E L P E R  F U N C T I O N S_______________________________________

float3 phongDistribution(float2 uv, float n, float3 axis)
{
	float theta = acos(pow(abs(1-uv.x), 1/(float)(n+1)));
	float phi = 2.0f*pi*uv.y;

	// (theta, phi) -> cartesian
	float3 res = float3(	sin(theta) * cos(phi),
							sin(theta) * sin(phi),
							cos(theta));

	float3 w = axis;
	w = normalize(w);

	float3 t = w;
	float t_min = min(t.x, t.y);
	t_min = min(t_min, t.z);
	if (t_min == t.x) t.x = 1.0f;
	else if (t_min == t.y) t.y = 1.0f;
	else if (t_min == t.z) t.z = 1.0;

	float3 u = cross(t,w);
	u = normalize(u);

	float3 v = cross(w,u);
	v = normalize(v);

	float3x3 xform = { u, v, w };
	res = mul(res, xform);
	return res;
}


float3 diskDistribution(float2 uv, float R, float3 normal)
{
	uv *= 2.0f;

	float theta = 2.0f * pi * uv.x;
	float r = R*sqrt(uv.y);
	r = R * sqrt(abs(uv.y));

	float x = r*cos(theta);
	float y = r*sin(theta);
	

	float3 Y = normal;
	Y = abs(Y);
	if (Y.x < Y.y && Y.x < Y.z)
		Y = float3(0, -Y.z, Y.y);
	else if (Y.y < Y.x && Y.y < Y.z)
		Y = float3(-Y.z, 0, Y.x);
	else 
		Y = float3(-Y.y, Y.x, 0);



	float3 X = normalize(cross(Y, normal));

	float3x3 mtx = { X, Y, normal };
	float3 res = mul(float3(x,y,0), mtx);
	return res;
}


//________________S H A D E R S_______________________________________
FireParticle_ID_GS VS_UpdateParticles(FireParticle_ID_VS vin)
{
	FireParticle_ID_GS vout = (FireParticle_ID_GS)0;
	vout.position =	 vin.position;
	vout.direction = vin.direction;
	vout.params =    vin.params;
	vout.vertexId =  vin.vertexId;

	if (vin.params.w >= 0.0f) // Emitters only
		vout.emitterWS = mul(float4(vin.position.xyz, 1.0f), matrices[vin.params.x]).xyz;

	return vout;
}


struct GS_TO_PS
{
	float4 pos : SV_POSITION;
};


//_________________________________________
[maxvertexcount(14)]
void GS_UpdateParticles(point FireParticle_ID_GS gin[1], inout PointStream<FireParticle> gstream)
														 /*,inout PointStream<GS_OUT>		 rasterizerStream)*/
{
	FireParticle particle;
	particle.position = gin[0].position;
	particle.direction = gin[0].direction;
	particle.params = gin[0].params;

	// check particle's type
	if (particle.params.w >= 0.0f) // emitters have params.w = ID = 0,1,2,..., for fire particles: params.w = -1 
	{
		// read heat value for this emitter from heat texture
		// if heat > threshold or if the activate_emitter flag is set to this 
		// emit new fire particles
		int x = gin[0].vertexId / heatTex_dimension;
		int y = gin[0].vertexId % heatTex_dimension;

		// update age
		particle.params.z += dt;

		if (gin[0].emitterWS.y < fHeightThreshold && dt - 0.0000001f > 0.0f)
		{
			// emit new particles
			FireParticle newP;
			// initial position equals emitter's position
			newP.position = particle.position;
			newP.position = float4(gin[0].emitterWS, 1.0f);

			// direction is the emitter's direction slightly altered 
			float3 rnd = tex_noise.SampleLevel(s_linear, float2(time*2.0f, -time), 0).rgb;
		
			float3 v = phongDistribution(rnd.xy, 21, particle.direction.xyz);
			newP.direction = normalize(float4(v,0));

			float minWindEffect = 0.4f;
			float maxWindEffect = 1.5f;
			float windEffect = minWindEffect + (maxWindEffect-minWindEffect)*fWindPower;

			if (bUseDiskDistribution) {
				v = diskDistribution(rnd.xy, 0.1, particle.direction.xyz);
				newP.position = float4(newP.position.xyz + v,1.0f);
				newP.direction = float4(0,1,0,0); 
			}

			// params:
			newP.params = float4(/*age*/fParticleAge, /*alpha*/1.0f , 1, /*type*/-1);
			gstream.Append(newP);

			particle.params.z = 0.0f; // age
			particle.params.y = 1.0f; // active
		}

		// always append emitter
		gstream.Append(particle);

	}
	else
	{
		// update position
		particle.position += particle.direction * dt;
		// update age
		particle.params.x = particle.params.x - dt;
		// update alpha value
		particle.params.y = particle.params.y - dt;

		float3 windDir = float3(vWindDir.x, 0.5f, vWindDir.z);
		windDir = normalize(windDir);
		particle.direction = float4(lerp(float3(0,1,0),windDir, fWindPower), 0);
		particle.direction.xyz = normalize(particle.direction.xyz);

		float minPower=1.0f, maxPower = 2.0f;
		particle.direction *= lerp(minPower, maxPower, fWindPower);

		if (particle.params.x > 0.0f)
			gstream.Append(particle);
	}
}

//_________________________________________
float4 PS_UpdateParticles(GS_OUT pin) : SV_TARGET
{
	return float4(1,1,1,1);
}


GeometryShader pGS = CompileShader(gs_4_0, GS_UpdateParticles());



GeometryShader gsUpdateParticles = 
			ConstructGSWithSO( 
								pGS, 
								"POSITION.xyzw; DIRECTION.xyzw; PARAMETERS.xyzw"
							);


DepthStencilState disabled
{
	DepthEnable = false;
	DepthWriteMask = false;
};



technique10 UpdateParticles
{
	pass P0
	{
		SetVertexShader( CompileShader ( vs_4_0, VS_UpdateParticles() ) );
		SetGeometryShader( gsUpdateParticles );
		SetPixelShader( NULL );
		
		SetDepthStencilState(disabled, 0);
	}
}

